import matplotlib.pyplot as plt
import numpy as np

"""
Dx/dt = - Y - Z
DY/dt = X + aY 
DZ/dt = b + Z( X - c)
"""
#指定参数
sigma = 0.5
gamma = 2.0
beta = 4.0
t = [0]
x = [1]
y = [0]
z = [0]
h = 0.01
N = 20000

def Rossler(x0,y0,z0,w,p,q,r,T):
  h=0.01
  x=[]
  y=[]
  z=[]
  for t in range(T):
    xt=x0+h*(-w*y0-z0)
    yt=y0+h*(w*x0+p*y0)
    zt=z0+h*(q+z0*(x0-r))

    #x0、y0、z0统一更新
    x0,y0,z0=xt,yt,zt
    x.append(x0)
    y.append(y0)
    z.append(z0)
  return x,y,z

#Dx/dt = - Y - Z
def FX(X, Y, Z):
    return - Y - Z

#DY/dt = X + aY
def FY(X, Y, Z):
    return X + sigma * Y

#DZ/dt = b + Z( X - c)
def FZ(X, Y, Z):
    return gamma + Z * (X - beta)

#龙格库塔方程组
def main():
    for i in range(N):
        K1 = FX(x[-1], y[-1], z[-1])
        L1 = FY(x[-1], y[-1], z[-1])
        M1 = FZ(x[-1], y[-1], z[-1])
        K2 = FX(x[-1] + h * K1 / 2, y[-1] + h * L1 / 2, z[-1] + h * M1 / 2)
        L2 = FY(x[-1] + h * K1 / 2, y[-1] + h * L1 / 2, z[-1] + h * M1 / 2)
        M2 = FZ(x[-1] + h * K1 / 2, y[-1] + h * L1 / 2, z[-1] + h * M1 / 2)
        K3 = FX(x[-1] + h * K2 / 2, y[-1] + h * L2 / 2, z[-1] + h * M2 / 2)
        L3 = FY(x[-1] + h * K2 / 2, y[-1] + h * L2 / 2, z[-1] + h * M2 / 2)
        M3 = FZ(x[-1] + h * K2 / 2, y[-1] + h * L2 / 2, z[-1] + h * M2 / 2)
        K4 = FX(x[-1] + h * K3, y[-1] + h * L3, z[-1] + h * M3)
        L4 = FY(x[-1] + h * K3, y[-1] + h * L3, z[-1] + h * M3)
        M4 = FZ(x[-1] + h * K3, y[-1] + h * L3, z[-1] + h * M3)
        x.append(x[-1] + h / 6 * (K1 + 2 * K2 + 2 * K3 + K4))
        y.append(y[-1] + h / 6 * (L1 + 2 * L2 + 2 * L3 + L4))
        z.append(z[-1] + h / 6 * (M1 + 2 * M2 + 2 * M3 + M4))
    #画三维图
    fig = plt.figure()
    ax = fig.add_subplot(projection='3d')
    ax.plot(x, y, z, label='parametric curve')
    plt.show()

    #设置参量
    w = 1
    p = 0.165
    q = 0.2
    #设定初值
    x0 = 1
    y0 = 0
    z0 = 0
    x1,y1,z1 = Rossler(x0, y0, z0, w, p, q, 10, N)
    #初值微小的变化
    x2 = 1
    y2 = 0
    z2 = 0.00001
    xx,yy,zz = Rossler(x2, y2, z2, w, p, q, 10, N)
    t = np.arange(0, N)
    plt.scatter(t, x1, s=1)
    plt.scatter(t, xx, s=1)
    plt.show()

if __name__=='__main__':
  main()